package bjut.zrpc.core.optimization.protect;


import java.net.InetSocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author: zhuzhiming
 * @Date: 2022/6/28 22:02
 */
public class ServicesRecord {

    public static ConcurrentHashMap<String, List<String>> SERVICES_MAP = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, List<String>> UNHEALTHY_SERVICE_MAP = new ConcurrentHashMap<>();

    private static final ConcurrentHashMap<String, Map<Long, Integer>> SERVICE_EXCEPTION_MAP = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, Long> SENTINEL_MAP = new ConcurrentHashMap<>();

    public void updateServiceSet(String serviceName, List<String> serviceList) {
        SERVICES_MAP.put(serviceName, serviceList);
        for (String service : SERVICES_MAP.keySet()) {
            UNHEALTHY_SERVICE_MAP.remove(service);
        }
    }

    public void addUnhealthyService(String serviceName, InetSocketAddress inetSocketAddress, boolean isTravel) {
        String service = inetSocketAddress.toString().replace("/", "");
        if (!isTravel) {
            List<String> serviceList = SERVICES_MAP.get(serviceName);
            if (serviceList != null) {
                serviceList.remove(service);
            }
        }
        List<String> unhealthyList = UNHEALTHY_SERVICE_MAP.get(serviceName);
        if (unhealthyList == null) {
            UNHEALTHY_SERVICE_MAP.put(serviceName, new ArrayList<>());
            unhealthyList = UNHEALTHY_SERVICE_MAP.get(serviceName);
        }
        unhealthyList.add(service);

    }

    public void removeFirst(String serviceName) {
        SERVICES_MAP.remove(serviceName);
    }


    public boolean needSentinel(String serviceNameAndAddress, int exceptionCount, long inTime) {

        String[] serviceNameAndAddressArr = serviceNameAndAddress.split("-");
        String serviceName = serviceNameAndAddressArr[0];
        String inetSocketAddress = serviceNameAndAddressArr[1];
        List<String> serviceList = SERVICES_MAP.get(serviceName);
        if (serviceList != null && serviceList.contains(inetSocketAddress)) {
            serviceList.remove(inetSocketAddress);
            return true;
        }

        Long window = SENTINEL_MAP.get(serviceNameAndAddress);
        if (window == null) {
            window = System.currentTimeMillis();
            SENTINEL_MAP.put(serviceNameAndAddress, window);
            HashMap<Long, Integer> exceptionMap = new HashMap<>();
            exceptionMap.put(window, 1);
            SERVICE_EXCEPTION_MAP.put(serviceNameAndAddress, exceptionMap);
        }

        long currentTime = System.currentTimeMillis();
        int currentCount = SERVICE_EXCEPTION_MAP.get(serviceNameAndAddress).get(window) + 1;

        if (currentTime < (window + inTime) && currentCount >= exceptionCount) {
            String[] inetSocket = inetSocketAddress.split(":");
            addUnhealthyService(serviceName, new InetSocketAddress(inetSocket[0], Integer.parseInt(inetSocket[1])), false);
            return true;
        }
        if (currentTime > (window + inTime) && currentCount < exceptionCount) {
            SENTINEL_MAP.remove(serviceNameAndAddress);
            SERVICE_EXCEPTION_MAP.get(serviceNameAndAddress).remove(window);
        }

        if (currentTime < (window + inTime) && currentCount < exceptionCount) {
            SERVICE_EXCEPTION_MAP.get(serviceNameAndAddress).put(window, currentCount);
        }

        return false;

    }

    public void checkSentinel(String serviceName, int duration) {

        long currentTime = System.currentTimeMillis();
        for (String key : SENTINEL_MAP.keySet()) {
            if (key.startsWith(serviceName)) {
                Long window = SENTINEL_MAP.get(key);
                if (currentTime > window + duration) {
                    List<String> serviceList = SERVICES_MAP.get(serviceName);
                    if (serviceList == null) {
                        SERVICES_MAP.put(serviceName, new ArrayList<>());
                        serviceList = SERVICES_MAP.get(serviceName);
                    }
                    String[] serviceNameAndAddressArr = key.split("-");
                    String serviceAddress = serviceNameAndAddressArr[1];
                    serviceList.add(serviceAddress);
                }
            }
        }

    }

}
